#!/usr/bin/env python3
"""
督导系统数据下载脚本
提供多种方式获取训练数据
"""

import os
import sys
import argparse
import logging
from datetime import datetime

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def download_from_urls():
    """从URL下载公开数据集（示例）"""
    
    logger.info("正在下载公开数据集...")
    
    # 这里可以添加真实的数据集下载链接
    # 由于督导数据比较特殊，我们使用生成的方式
    
    logger.warning("目前没有可用的公开督导数据集")
    logger.info("建议使用生成器创建训练数据")
    
    return False

def generate_synthetic_data(size: str = "medium"):
    """生成合成训练数据"""
    
    logger.info(f"生成 {size} 规模的合成训练数据...")
    
    # 导入数据生成器
    from data_generator import SupervisionDataGenerator, DataGeneratorConfig
    
    # 根据规模设置参数
    size_configs = {
        "small": {"num_areas": 50, "num_records": 5000, "time_span_days": 365},
        "medium": {"num_areas": 150, "num_records": 20000, "time_span_days": 900},
        "large": {"num_areas": 300, "num_records": 50000, "time_span_days": 1095},
        "xlarge": {"num_areas": 500, "num_records": 100000, "time_span_days": 1460}
    }
    
    if size not in size_configs:
        logger.error(f"不支持的数据规模: {size}")
        return False
    
    config = DataGeneratorConfig(**size_configs[size])
    generator = SupervisionDataGenerator(config)
    
    try:
        # 生成数据
        areas_df, records_df = generator.generate_dataset()
        
        # 保存数据
        generator.save_dataset(areas_df, records_df)
        
        # 生成SQL脚本
        from data_generator import create_database_sql
        create_database_sql(areas_df, records_df)
        
        logger.info("✅ 数据生成完成！")
        return True
        
    except Exception as e:
        logger.error(f"数据生成失败: {e}")
        return False

def check_existing_data():
    """检查现有数据"""
    
    logger.info("检查现有数据文件...")
    
    data_files = [
        "data/areas.csv",
        "data/supervision_records.csv",
        "data/dataset_stats.json"
    ]
    
    existing_files = []
    for file_path in data_files:
        if os.path.exists(file_path):
            size = os.path.getsize(file_path)
            existing_files.append((file_path, size))
            logger.info(f"✅ {file_path} ({size:,} bytes)")
        else:
            logger.info(f"❌ {file_path} (不存在)")
    
    if existing_files:
        logger.info(f"找到 {len(existing_files)} 个现有数据文件")
        
        # 读取统计信息
        stats_file = "data/dataset_stats.json"
        if os.path.exists(stats_file):
            import json
            with open(stats_file, 'r', encoding='utf-8') as f:
                stats = json.load(f)
            
            logger.info("📊 数据集统计信息:")
            logger.info(f"  区域数量: {stats['statistics']['areas_count']}")
            logger.info(f"  督导记录数量: {stats['statistics']['records_count']}")
            logger.info(f"  时间范围: {stats['statistics']['date_range']['start']} 到 {stats['statistics']['date_range']['end']}")
        
        return True
    else:
        logger.info("未找到现有数据文件")
        return False

def setup_database():
    """设置数据库（可选）"""
    
    logger.info("数据库设置说明:")
    logger.info("1. 如果您想使用数据库，请先创建PostgreSQL数据库")
    logger.info("2. 在项目根目录创建.env文件，配置数据库连接:")
    logger.info("   DB_HOST=localhost")
    logger.info("   DB_PORT=5432") 
    logger.info("   DB_NAME=supervision_db")
    logger.info("   DB_USER=postgres")
    logger.info("   DB_PASSWORD=your_password")
    logger.info("3. 运行SQL脚本创建表和导入数据:")
    logger.info("   psql -U postgres -d supervision_db -f data/create_database.sql")
    logger.info("\n注意: 如果不使用数据库，模型训练会自动从CSV文件加载数据")

def show_usage_examples():
    """显示使用示例"""
    
    logger.info("📖 使用示例:")
    logger.info("")
    logger.info("1. 生成小规模数据集进行快速测试:")
    logger.info("   python download_dataset.py --generate --size small")
    logger.info("")
    logger.info("2. 生成中等规模数据集进行正式训练:")
    logger.info("   python download_dataset.py --generate --size medium")
    logger.info("")
    logger.info("3. 生成大规模数据集:")
    logger.info("   python download_dataset.py --generate --size large")
    logger.info("")
    logger.info("4. 检查现有数据:")
    logger.info("   python download_dataset.py --check")
    logger.info("")
    logger.info("5. 生成数据后开始训练:")
    logger.info("   python train_models.py --models all")

def main():
    """主函数"""
    
    parser = argparse.ArgumentParser(description='督导系统训练数据下载器')
    
    parser.add_argument('--generate', action='store_true', 
                       help='生成合成训练数据')
    parser.add_argument('--size', choices=['small', 'medium', 'large', 'xlarge'],
                       default='medium', help='数据集规模')
    parser.add_argument('--check', action='store_true',
                       help='检查现有数据文件')
    parser.add_argument('--setup-db', action='store_true',
                       help='显示数据库设置说明')
    parser.add_argument('--examples', action='store_true',
                       help='显示使用示例')
    
    args = parser.parse_args()
    
    # 如果没有参数，显示帮助
    if len(sys.argv) == 1:
        parser.print_help()
        logger.info("\n" + "="*50)
        show_usage_examples()
        return
    
    logger.info("🚀 督导系统数据下载器")
    logger.info("="*50)
    
    success = True
    
    if args.check:
        check_existing_data()
    
    if args.generate:
        success = generate_synthetic_data(args.size)
        
        if success:
            logger.info("\n" + "="*50)
            logger.info("🎉 数据生成成功！")
            logger.info("\n📁 生成的文件:")
            logger.info("  - data/areas.csv (区域信息)")
            logger.info("  - data/supervision_records.csv (督导记录)")
            logger.info("  - data/dataset_stats.json (统计信息)")
            logger.info("  - data/create_database.sql (数据库脚本)")
            logger.info("\n🚀 下一步:")
            logger.info("  python train_models.py --models all")
    
    if args.setup_db:
        setup_database()
    
    if args.examples:
        show_usage_examples()
    
    if args.generate and not success:
        logger.error("❌ 数据生成失败")
        sys.exit(1)

if __name__ == "__main__":
    main() 